import numpy as np
import torch
import torch.nn as nn
from algorithms.utils.util import huber_loss
from utils.valuenorm import ValueNorm
from algorithms.utils.util import check


# MAT训练器类，用于更新策略。
class MATTrainer:

    def __init__(self,
                 args,
                 policy,
                 num_agents,
                 device=torch.device("cpu")):
        """
        初始化MAT训练器
        :param args: 训练参数配置
        :param policy: 策略网络（TransformerPolicy）
        :param num_agents: 智能体数量
        :param device: 计算设备（CPU或GPU）
        """
        self.device = device  # 计算设备
        self.tpdv = dict(dtype=torch.float32, device=device)  # 张量参数字典
        self.policy = policy  # 策略网络
        self.num_agents = num_agents  # 智能体数量

        # PPO算法相关参数
        self.clip_param = args.clip_param  # PPO裁剪参数，用于限制策略更新幅度
        self.ppo_epoch = args.ppo_epoch  # PPO更新轮数，每轮使用相同数据多次更新
        self.num_mini_batch = args.num_mini_batch  # 小批量数量，用于数据分批处理
        self.value_loss_coef = args.value_loss_coef  # 价值损失系数，控制价值函数损失在总损失中的权重
        self.entropy_coef = args.entropy_coef  # 熵损失系数，鼓励探索
        self.max_grad_norm = args.max_grad_norm  # 最大梯度范数，用于梯度裁剪
        self.huber_delta = args.huber_delta  # Huber损失参数，用于价值函数损失计算

        # 初始化价值归一化器
        self.value_normalizer = ValueNorm(1, device=self.device)  # 价值归一化器

    def cal_value_loss(self, values, value_preds_batch, return_batch, active_masks_batch):
        """
        计算价值函数损失
        :param values: (torch.Tensor) 当前价值函数预测值
        :param value_preds_batch: (torch.Tensor) 数据批次中的"旧"价值预测（用于价值裁剪损失）
        :param return_batch: (torch.Tensor) 奖励回报（真实价值）
        :param active_masks_batch: (torch.Tensor) 活跃掩码，表示智能体在给定时间步是否活跃或死亡
        :return value_loss: (torch.Tensor) 计算得到的价值函数损失
        """
        # 价值函数裁剪：限制价值函数更新幅度，防止过度更新
        value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param,
                                                                                    self.clip_param)

        # 价值归一化处理
        self.value_normalizer.update(return_batch)  # 更新归一化器统计信息
        error_clipped = self.value_normalizer.normalize(return_batch) - value_pred_clipped  # 裁剪后的误差
        error_original = self.value_normalizer.normalize(return_batch) - values  # 原始误差

        # Huber损失
        value_loss_clipped = huber_loss(error_clipped, self.huber_delta)
        value_loss_original = huber_loss(error_original, self.huber_delta)

        # 价值损失裁剪：取原始损失和裁剪损失的最大值
        value_loss = torch.max(value_loss_original, value_loss_clipped)

        # 活跃掩码处理：只计算活跃智能体的损失
        value_loss = (value_loss * active_masks_batch).sum() / active_masks_batch.sum()

        return value_loss

    def ppo_update(self, sample):
        """
        执行一次PPO更新，更新actor和critic网络
        :param sample: (Tuple) 包含用于更新网络的数据批次，包括：
                      - share_obs_batch: 共享观测
                      - obs_batch: 局部观测
                      - actions_batch: 动作
                      - value_preds_batch: 价值预测
                      - return_batch: 回报
                      - active_masks_batch: 活跃掩码
                      - old_action_log_probs_batch: 旧动作对数概率
                      - adv_targ: 优势函数
                      - available_actions_batch: 可用动作
        :return value_loss: (torch.Tensor) 价值函数损失
        :return critic_grad_norm: (torch.Tensor) critic更新的梯度范数
        :return policy_loss: (torch.Tensor) actor(策略)损失值
        :return dist_entropy: (torch.Tensor) 动作熵
        :return actor_grad_norm: (torch.Tensor) actor更新的梯度范数
        :return imp_weights: (torch.Tensor) 重要性采样权重
        """
        # 解包数据批次
        obs_batch, actions_batch, \
            value_preds_batch, return_batch, active_masks_batch, old_action_log_probs_batch, \
            adv_targ, available_actions_batch = sample

        # 将数据转换为张量并移动到指定设备
        old_action_log_probs_batch = check(old_action_log_probs_batch).to(**self.tpdv)  # 旧动作对数概率
        adv_targ = check(adv_targ).to(**self.tpdv)  # 优势函数
        value_preds_batch = check(value_preds_batch).to(**self.tpdv)  # 价值预测
        return_batch = check(return_batch).to(**self.tpdv)  # 回报
        active_masks_batch = check(active_masks_batch).to(**self.tpdv)  # 活跃掩码

        # 策略评估：计算当前策略下的价值、动作对数概率和熵
        values, action_log_probs, dist_entropy = self.policy.evaluate_actions(obs_batch,
                                                                              actions_batch,
                                                                              available_actions_batch,
                                                                              active_masks_batch)

        # Actor更新：计算重要性采样权重
        imp_weights = torch.exp(action_log_probs - old_action_log_probs_batch)  # 重要性采样权重

        # PPO裁剪目标函数：计算两个替代目标
        surr1 = imp_weights * adv_targ  # 第一个替代目标：重要性采样权重 * 优势
        surr2 = torch.clamp(imp_weights, 1.0 - self.clip_param,
                            1.0 + self.clip_param) * adv_targ  # 第二个替代目标：裁剪后的重要性采样权重 * 优势

        # 策略损失计算：取两个替代目标的最小值，并考虑活跃掩码
        policy_loss = (-torch.sum(torch.min(surr1, surr2),
                                  dim=-1,
                                  keepdim=True) * active_masks_batch).sum() / active_masks_batch.sum()

        # Critic更新：计算价值函数损失
        value_loss = self.cal_value_loss(values, value_preds_batch, return_batch, active_masks_batch)

        # 总损失：策略损失 - 熵损失 + 价值损失
        loss = policy_loss - dist_entropy * self.entropy_coef + value_loss * self.value_loss_coef

        # 梯度更新
        self.policy.optimizer.zero_grad()  # 清空梯度
        loss.backward()  # 反向传播

        # 梯度裁剪：防止梯度爆炸
        grad_norm = nn.utils.clip_grad_norm_(self.policy.transformer.parameters(), self.max_grad_norm)

        self.policy.optimizer.step()  # 参数更新

        return value_loss, grad_norm, policy_loss, dist_entropy, grad_norm, imp_weights

    def train(self, buffer):
        """
        使用小批量梯度下降执行训练更新
        :param buffer: (SharedReplayBuffer) 包含训练数据的缓冲区
        :return train_info: (dict) 包含训练更新相关信息（例如损失、梯度范数等）
        """
        # 优势函数标准化：提高训练稳定性
        advantages_copy = buffer.advantages.copy()  # 复制优势函数
        advantages_copy[buffer.active_masks[:-1] == 0.0] = np.nan  # 将不活跃的智能体优势设为NaN
        mean_advantages = np.nanmean(advantages_copy)  # 计算平均优势
        std_advantages = np.nanstd(advantages_copy)  # 计算优势标准差
        advantages = (buffer.advantages - mean_advantages) / (std_advantages + 1e-5)  # 标准化优势函数

        # 初始化训练信息字典
        train_info = {}

        train_info['value_loss'] = 0  # 价值损失
        train_info['policy_loss'] = 0  # 策略损失
        train_info['dist_entropy'] = 0  # 分布熵
        train_info['actor_grad_norm'] = 0  # Actor梯度范数
        train_info['critic_grad_norm'] = 0  # Critic梯度范数
        train_info['ratio'] = 0  # 重要性采样权重平均值

        # 多轮PPO更新：使用相同数据多次更新策略
        for _ in range(self.ppo_epoch):
            # 数据生成器：将数据分成小批量
            data_generator = buffer.feed_forward_generator_transformer(advantages, self.num_mini_batch)

            # 遍历每个小批量
            for sample in data_generator:
                # 执行PPO更新
                value_loss, critic_grad_norm, policy_loss, dist_entropy, actor_grad_norm, imp_weights \
                    = self.ppo_update(sample)

                # 累积训练信息
                train_info['value_loss'] += value_loss.item()
                train_info['policy_loss'] += policy_loss.item()
                train_info['dist_entropy'] += dist_entropy.item()
                train_info['actor_grad_norm'] += actor_grad_norm
                train_info['critic_grad_norm'] += critic_grad_norm
                train_info['ratio'] += imp_weights.mean()

        # 计算平均训练信息
        num_updates = self.ppo_epoch * self.num_mini_batch  # 总更新次数

        for k in train_info.keys():
            train_info[k] /= num_updates  # 计算平均值

        return train_info

    def prep_training(self):
        """
        准备训练模式：将策略网络设置为训练模式
        """
        self.policy.train()

    def prep_rollout(self):
        """
        准备回滚模式：将策略网络设置为评估模式
        """
        self.policy.eval()

    def get_actions(self, obs, available_actions=None, deterministic=False):
        """
        获取动作 - 用于数据收集
        
        Args:
            obs: 观测张量
            available_actions: 可用动作掩码
            deterministic: 是否确定性采样
            
        Returns:
            values: 价值函数输出
            actions: 选择的动作
            action_log_probs: 动作对数概率
        """
        # 准备策略网络为推理模式
        self.prep_rollout()
        
        # 使用策略网络获取动作
        actions, action_log_probs, values = self.policy.get_actions(
            obs, available_actions, deterministic
        )
        
        return values, actions, action_log_probs


# 4D观测MAT端到端训练器
class MATTrainerSimple(nn.Module):
    """
    简化的4D训练器，用于快速原型验证
    """

    def __init__(self, policy, device=torch.device("cpu")):
        """
        初始化简化训练器
        
        Args:
            policy: 4D策略
            device: 计算设备
        """
        super(MATTrainerSimple, self).__init__()
        self.policy = policy
        self.device = device
        self.tpdv = dict(dtype=torch.float32, device=device)
        
        print(f"简化4D训练器初始化完成，设备: {device}")

    def simple_train_step(self, obs_batch, actions_batch, returns_batch, advantages_batch):
        """
        简化的训练步骤
        
        Args:
            obs_batch: 观测批次 (batch_size, n_agents, 36)
            actions_batch: 动作批次 (batch_size, n_agents, 1)
            returns_batch: 回报批次 (batch_size, n_agents, 1)
            advantages_batch: 优势批次 (batch_size, n_agents, 1)
            
        Returns:
            loss: 训练损失
        """
        # 转换为torch张量
        obs_tensor = torch.from_numpy(obs_batch).to(**self.tpdv)
        actions_tensor = torch.from_numpy(actions_batch).to(**self.tpdv)
        returns_tensor = torch.from_numpy(returns_batch).to(**self.tpdv)
        advantages_tensor = torch.from_numpy(advantages_batch).to(**self.tpdv)
        
        # 评估动作 - 使用evaluate_actions获取一致的价值和动作概率
        values_eval, action_log_probs_eval, dist_entropy = self.policy.evaluate_actions(
            obs_batch, actions_batch
        )
        
        # 计算损失 - 使用evaluate_actions返回的values，确保维度匹配
        # values_eval: (batch_size, n_agents, 1), returns_tensor: (batch_size, n_agents)
        # 需要将returns_tensor扩展一个维度以匹配values_eval
        returns_tensor_expanded = returns_tensor.unsqueeze(-1)  # (batch_size, n_agents, 1)
        value_loss = torch.mean((values_eval - returns_tensor_expanded) ** 2)
        
        # 确保action_log_probs_eval和advantages_tensor维度匹配
        # action_log_probs_eval: (batch_size, n_agents, 1), advantages_tensor: (batch_size, n_agents)
        advantages_tensor_expanded = advantages_tensor.unsqueeze(-1)  # (batch_size, n_agents, 1)
        policy_loss = -torch.mean(action_log_probs_eval * advantages_tensor_expanded)
        entropy_loss = -torch.mean(dist_entropy)
        
        # 总损失
        total_loss = value_loss + policy_loss + 0.01 * entropy_loss
        
        # 更新
        self.policy.update(total_loss)
        
        return {
            'total_loss': total_loss.item(),
            'value_loss': value_loss.item(),
            'policy_loss': policy_loss.item(),
            'entropy_loss': entropy_loss.item(),
        }
